{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the model on Taiwan Street Accident datasets\n",
    "* split data to train/test\n",
    "* train data only has normal, test data has normal and abnormal.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import os\n",
    "\n",
    "from models import generator, discriminator, flownet, initialize_flownet\n",
    "from loss_functions import intensity_loss, gradient_loss\n",
    "from utils import DataLoader, load, save, psnr_error\n",
    "from constant import const\n",
    "import tensorboard\n",
    "\n",
    "\n",
    "os.environ['CUDA_DEVICES_ORDER'] = \"PCI_BUS_ID\"\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = const.GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DATASET taiwan_sa\n",
      "TRAIN_FOLDER /media/DATA/VAD_datasets/taiwan_sa/training/frames\n",
      "TEST_FOLDER /media/DATA/VAD_datasets/taiwan_sa/testing/frames\n",
      "GPU 0\n",
      "BATCH_SIZE 4\n",
      "NUM_HIS 4\n",
      "ITERATIONS 10000\n",
      "EVALUATE compute_auc\n",
      "HEIGHT 256\n",
      "WIDTH 256\n",
      "FLOWNET_CHECKPOINT checkpoints/pretrains/flownet-SD.ckpt-0\n",
      "FLOW_HEIGHT 384\n",
      "FLOW_WIDTH 512\n",
      "L_NUM 2\n",
      "ALPHA_NUM 1\n",
      "LAM_ADV 0.05\n",
      "LAM_LP 1.0\n",
      "LAM_GDL 1.0\n",
      "LAM_FLOW 2.0\n",
      "LRATE_G [0.0002, 2e-05]\n",
      "LRATE_G_BOUNDARIES [50000]\n",
      "LRATE_D [2e-05, 2e-06]\n",
      "LRATE_D_BOUNDARIES [50000]\n",
      "SAVE_DIR taiwan_sa_l_2_alpha_1_lp_1.0_adv_0.05_gdl_1.0_flow_2.0\n",
      "SNAPSHOT_DIR checkpoints/taiwan_sa_l_2_alpha_1_lp_1.0_adv_0.05_gdl_1.0_flow_2.0\n",
      "SUMMARY_DIR summary/taiwan_sa_l_2_alpha_1_lp_1.0_adv_0.05_gdl_1.0_flow_2.0\n",
      "PSNR_DIR psnrs/taiwan_sa_l_2_alpha_1_lp_1.0_adv_0.05_gdl_1.0_flow_2.0\n",
      "<================ Constants information ================>\n",
      "\tDATASET\ttaiwan_sa\n",
      "\tTRAIN_FOLDER\t/media/DATA/VAD_datasets/taiwan_sa/training/frames\n",
      "\tTEST_FOLDER\t/media/DATA/VAD_datasets/taiwan_sa/testing/frames\n",
      "\tGPU\t0\n",
      "\tBATCH_SIZE\t4\n",
      "\tNUM_HIS\t4\n",
      "\tITERATIONS\t10000\n",
      "\tEVALUATE\tcompute_auc\n",
      "\tHEIGHT\t256\n",
      "\tWIDTH\t256\n",
      "\tFLOWNET_CHECKPOINT\tcheckpoints/pretrains/flownet-SD.ckpt-0\n",
      "\tFLOW_HEIGHT\t384\n",
      "\tFLOW_WIDTH\t512\n",
      "\tL_NUM\t2\n",
      "\tALPHA_NUM\t1\n",
      "\tLAM_ADV\t0.05\n",
      "\tLAM_LP\t1.0\n",
      "\tLAM_GDL\t1.0\n",
      "\tLAM_FLOW\t2.0\n",
      "\tLRATE_G\t[0.0002, 2e-05]\n",
      "\tLRATE_G_BOUNDARIES\t[50000]\n",
      "\tLRATE_D\t[2e-05, 2e-06]\n",
      "\tLRATE_D_BOUNDARIES\t[50000]\n",
      "\tSAVE_DIR\ttaiwan_sa_l_2_alpha_1_lp_1.0_adv_0.05_gdl_1.0_flow_2.0\n",
      "\tSNAPSHOT_DIR\tcheckpoints/taiwan_sa_l_2_alpha_1_lp_1.0_adv_0.05_gdl_1.0_flow_2.0\n",
      "\tSUMMARY_DIR\tsummary/taiwan_sa_l_2_alpha_1_lp_1.0_adv_0.05_gdl_1.0_flow_2.0\n",
      "\tPSNR_DIR\tpsnrs/taiwan_sa_l_2_alpha_1_lp_1.0_adv_0.05_gdl_1.0_flow_2.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "dataset_name = const.DATASET\n",
    "train_folder = const.TRAIN_FOLDER\n",
    "test_folder = const.TEST_FOLDER\n",
    "\n",
    "batch_size = const.BATCH_SIZE\n",
    "iterations = const.ITERATIONS\n",
    "num_his = const.NUM_HIS\n",
    "height, width = 256, 256\n",
    "flow_height, flow_width = const.FLOW_HEIGHT, const.FLOW_WIDTH\n",
    "\n",
    "l_num = const.L_NUM\n",
    "alpha_num = const.ALPHA_NUM\n",
    "lam_lp = const.LAM_LP\n",
    "lam_gdl = const.LAM_GDL\n",
    "lam_adv = const.LAM_ADV\n",
    "lam_flow = const.LAM_FLOW\n",
    "adversarial = (lam_adv != 0)\n",
    "\n",
    "summary_dir = const.SUMMARY_DIR\n",
    "snapshot_dir = const.SNAPSHOT_DIR\n",
    "\n",
    "\n",
    "print(const)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "generator dataset, <FlatMapDataset shapes: (256, 256, 15), types: tf.float32>\n",
      "epoch dataset, <BatchDataset shapes: (?, 256, 256, 15), types: tf.float32>\n",
      "train inputs = Tensor(\"dataset/strided_slice:0\", shape=(4, 256, 256, 12), dtype=float32)\n",
      "train prediction gt = Tensor(\"dataset/strided_slice_1:0\", shape=(4, 256, 256, 3), dtype=float32)\n",
      "generator dataset, <FlatMapDataset shapes: (256, 256, 15), types: tf.float32>\n",
      "epoch dataset, <BatchDataset shapes: (?, 256, 256, 15), types: tf.float32>\n",
      "test inputs = Tensor(\"dataset/strided_slice_2:0\", shape=(4, 256, 256, 12), dtype=float32)\n",
      "test prediction gt = Tensor(\"dataset/strided_slice_3:0\", shape=(4, 256, 256, 3), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# define dataset\n",
    "with tf.name_scope('dataset'):\n",
    "    train_loader = DataLoader(train_folder, resize_height=height, resize_width=width)\n",
    "    train_dataset = train_loader(batch_size=batch_size, time_steps=num_his, num_pred=1)\n",
    "\n",
    "    train_it = train_dataset.make_one_shot_iterator()\n",
    "    train_videos_clips_tensor = train_it.get_next()\n",
    "    train_videos_clips_tensor.set_shape([batch_size, height, width, 3*(num_his + 1)])\n",
    "\n",
    "    train_inputs = train_videos_clips_tensor[..., 0:num_his*3]\n",
    "    train_gt = train_videos_clips_tensor[..., -3:]\n",
    "\n",
    "    print('train inputs = {}'.format(train_inputs))\n",
    "    print('train prediction gt = {}'.format(train_gt))\n",
    "\n",
    "    test_loader = DataLoader(test_folder, resize_height=height, resize_width=width)\n",
    "    test_dataset = test_loader(batch_size=batch_size, time_steps=num_his, num_pred=1)\n",
    "    test_it = test_dataset.make_one_shot_iterator()\n",
    "    test_videos_clips_tensor = test_it.get_next()\n",
    "    test_videos_clips_tensor.set_shape([batch_size, height, width, 3*(num_his + 1)])\n",
    "\n",
    "    test_inputs = test_videos_clips_tensor[..., 0:num_his*3]\n",
    "    test_gt = test_videos_clips_tensor[..., -3:]\n",
    "\n",
    "    print('test inputs = {}'.format(test_inputs))\n",
    "    print('test prediction gt = {}'.format(test_gt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training = generator\n",
      "testing = generator\n",
      "WARNING:tensorflow:From /home/brianyao/Documents/VAD_baselines/ano_pred_cvpr2018/Codes/flownet2/src/net.py:26: get_or_create_global_step (from tensorflow.contrib.framework.python.ops.variables) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please switch to tf.train.get_or_create_global_step\n",
      "real_outputs = Tensor(\"discriminator/Sigmoid:0\", shape=(4, 35, 35, 1), dtype=float32)\n",
      "fake_outputs = Tensor(\"discriminator_1/Sigmoid:0\", shape=(4, 35, 35, 1), dtype=float32)\n",
      "Init successfully!\n",
      "FlownetSD restore from checkpoints/pretrains/flownet-SD.ckpt-0!\n",
      "INFO:tensorflow:Restoring parameters from checkpoints/pretrains/flownet-SD.ckpt-0\n",
      "No checkpoint file found.\n",
      "10\n",
      "DiscriminatorModel: Step 10 | Global Loss: 0.222318, lr = 0.000020\n",
      "GeneratorModel : Step 10, lr = 0.000200\n",
      "                 Global      Loss :  4.251993\n",
      "                 intensity   Loss : (0.1292 * 1.0000 = 0.1292)\n",
      "                 gradient    Loss : (0.1342 * 1.0000 = 0.1342)\n",
      "                 adversarial Loss : (0.0908 * 0.0500 = 0.0045)\n",
      "                 flownet     Loss : (1.9920 * 2.0000 = 3.9841)\n",
      "                 PSNR  Error      :  14.992388\n",
      "20\n",
      "DiscriminatorModel: Step 20 | Global Loss: 0.205181, lr = 0.000020\n",
      "GeneratorModel : Step 20, lr = 0.000200\n",
      "                 Global      Loss :  4.950412\n",
      "                 intensity   Loss : (0.0868 * 1.0000 = 0.0868)\n",
      "                 gradient    Loss : (0.1220 * 1.0000 = 0.1220)\n",
      "                 adversarial Loss : (0.0923 * 0.0500 = 0.0046)\n",
      "                 flownet     Loss : (2.3685 * 2.0000 = 4.7370)\n",
      "                 PSNR  Error      :  16.888992\n",
      "30\n",
      "DiscriminatorModel: Step 30 | Global Loss: 0.173388, lr = 0.000020\n",
      "GeneratorModel : Step 30, lr = 0.000200\n",
      "                 Global      Loss :  2.969307\n",
      "                 intensity   Loss : (0.1594 * 1.0000 = 0.1594)\n",
      "                 gradient    Loss : (0.1241 * 1.0000 = 0.1241)\n",
      "                 adversarial Loss : (0.1131 * 0.0500 = 0.0057)\n",
      "                 flownet     Loss : (1.3401 * 2.0000 = 2.6802)\n",
      "                 PSNR  Error      :  14.379272\n",
      "40\n",
      "DiscriminatorModel: Step 40 | Global Loss: 0.144297, lr = 0.000020\n",
      "GeneratorModel : Step 40, lr = 0.000200\n",
      "                 Global      Loss :  3.3850996\n",
      "                 intensity   Loss : (0.0978 * 1.0000 = 0.0978)\n",
      "                 gradient    Loss : (0.1050 * 1.0000 = 0.1050)\n",
      "                 adversarial Loss : (0.1345 * 0.0500 = 0.0067)\n",
      "                 flownet     Loss : (1.5878 * 2.0000 = 3.1757)\n",
      "                 PSNR  Error      :  16.352108\n",
      "50\n",
      "DiscriminatorModel: Step 50 | Global Loss: 0.090885, lr = 0.000020\n",
      "GeneratorModel : Step 50, lr = 0.000200\n",
      "                 Global      Loss :  3.3079991\n",
      "                 intensity   Loss : (0.0802 * 1.0000 = 0.0802)\n",
      "                 gradient    Loss : (0.1447 * 1.0000 = 0.1447)\n",
      "                 adversarial Loss : (0.2098 * 0.0500 = 0.0105)\n",
      "                 flownet     Loss : (1.5363 * 2.0000 = 3.0726)\n",
      "                 PSNR  Error      :  17.297659\n",
      "60\n",
      "DiscriminatorModel: Step 60 | Global Loss: 0.174095, lr = 0.000020\n",
      "GeneratorModel : Step 60, lr = 0.000200\n",
      "                 Global      Loss :  2.2345307\n",
      "                 intensity   Loss : (0.0825 * 1.0000 = 0.0825)\n",
      "                 gradient    Loss : (0.1190 * 1.0000 = 0.1190)\n",
      "                 adversarial Loss : (0.1982 * 0.0500 = 0.0099)\n",
      "                 flownet     Loss : (1.0115 * 2.0000 = 2.0230)\n",
      "                 PSNR  Error      :  16.911442\n",
      "70\n",
      "DiscriminatorModel: Step 70 | Global Loss: 0.367785, lr = 0.000020\n",
      "GeneratorModel : Step 70, lr = 0.000200\n",
      "                 Global      Loss :  4.178986\n",
      "                 intensity   Loss : (0.1230 * 1.0000 = 0.1230)\n",
      "                 gradient    Loss : (0.1645 * 1.0000 = 0.1645)\n",
      "                 adversarial Loss : (0.2786 * 0.0500 = 0.0139)\n",
      "                 flownet     Loss : (1.9388 * 2.0000 = 3.8776)\n",
      "                 PSNR  Error      :  15.436696\n",
      "80\n",
      "DiscriminatorModel: Step 80 | Global Loss: 0.348793, lr = 0.000020\n",
      "GeneratorModel : Step 80, lr = 0.000200\n",
      "                 Global      Loss :  4.1084723\n",
      "                 intensity   Loss : (0.0740 * 1.0000 = 0.0740)\n",
      "                 gradient    Loss : (0.1154 * 1.0000 = 0.1154)\n",
      "                 adversarial Loss : (0.3384 * 0.0500 = 0.0169)\n",
      "                 flownet     Loss : (1.9511 * 2.0000 = 3.9021)\n",
      "                 PSNR  Error      :  17.443777\n",
      "90\n",
      "DiscriminatorModel: Step 90 | Global Loss: 0.220073, lr = 0.000020\n",
      "GeneratorModel : Step 90, lr = 0.000200\n",
      "                 Global      Loss :  1.645825\n",
      "                 intensity   Loss : (0.0648 * 1.0000 = 0.0648)\n",
      "                 gradient    Loss : (0.1231 * 1.0000 = 0.1231)\n",
      "                 adversarial Loss : (0.1957 * 0.0500 = 0.0098)\n",
      "                 flownet     Loss : (0.7241 * 2.0000 = 1.4482)\n",
      "                 PSNR  Error      :  18.032877\n",
      "100\n",
      "DiscriminatorModel: Step 100 | Global Loss: 0.168513, lr = 0.000020\n",
      "GeneratorModel : Step 100, lr = 0.000200\n",
      "                 Global      Loss :  3.8679967\n",
      "                 intensity   Loss : (0.0513 * 1.0000 = 0.0513)\n",
      "                 gradient    Loss : (0.1163 * 1.0000 = 0.1163)\n",
      "                 adversarial Loss : (0.1819 * 0.0500 = 0.0091)\n",
      "                 flownet     Loss : (1.8457 * 2.0000 = 3.6913)\n",
      "                 PSNR  Error      :  18.988342\n",
      "Save summaries...\n",
      "110\n",
      "DiscriminatorModel: Step 110 | Global Loss: 0.121466, lr = 0.000020\n",
      "GeneratorModel : Step 110, lr = 0.000200\n",
      "                 Global      Loss :  4.461573\n",
      "                 intensity   Loss : (0.1066 * 1.0000 = 0.1066)\n",
      "                 gradient    Loss : (0.1720 * 1.0000 = 0.1720)\n",
      "                 adversarial Loss : (0.2102 * 0.0500 = 0.0105)\n",
      "                 flownet     Loss : (2.0862 * 2.0000 = 4.1725)\n",
      "                 PSNR  Error      :  15.846859\n",
      "120\n",
      "DiscriminatorModel: Step 120 | Global Loss: 0.130620, lr = 0.000020\n",
      "GeneratorModel : Step 120, lr = 0.000200\n",
      "                 Global      Loss :  3.6381469\n",
      "                 intensity   Loss : (0.0796 * 1.0000 = 0.0796)\n",
      "                 gradient    Loss : (0.1563 * 1.0000 = 0.1563)\n",
      "                 adversarial Loss : (0.2289 * 0.0500 = 0.0114)\n",
      "                 flownet     Loss : (1.6954 * 2.0000 = 3.3908)\n",
      "                 PSNR  Error      :  17.271368\n",
      "130\n",
      "DiscriminatorModel: Step 130 | Global Loss: 0.161342, lr = 0.000020\n",
      "GeneratorModel : Step 130, lr = 0.000200\n",
      "                 Global      Loss :  3.9195824\n",
      "                 intensity   Loss : (0.0689 * 1.0000 = 0.0689)\n",
      "                 gradient    Loss : (0.1163 * 1.0000 = 0.1163)\n",
      "                 adversarial Loss : (0.2435 * 0.0500 = 0.0122)\n",
      "                 flownet     Loss : (1.8611 * 2.0000 = 3.7222)\n",
      "                 PSNR  Error      :  17.827011\n",
      "140\n",
      "DiscriminatorModel: Step 140 | Global Loss: 0.407582, lr = 0.000020\n",
      "GeneratorModel : Step 140, lr = 0.000200\n",
      "                 Global      Loss :  0.9506073\n",
      "                 intensity   Loss : (0.0666 * 1.0000 = 0.0666)\n",
      "                 gradient    Loss : (0.1695 * 1.0000 = 0.1695)\n",
      "                 adversarial Loss : (0.2294 * 0.0500 = 0.0115)\n",
      "                 flownet     Loss : (0.3515 * 2.0000 = 0.7030)\n",
      "                 PSNR  Error      :  18.011234\n",
      "150\n",
      "DiscriminatorModel: Step 150 | Global Loss: 0.232443, lr = 0.000020\n",
      "GeneratorModel : Step 150, lr = 0.000200\n",
      "                 Global      Loss :  3.8960745\n",
      "                 intensity   Loss : (0.1001 * 1.0000 = 0.1001)\n",
      "                 gradient    Loss : (0.1546 * 1.0000 = 0.1546)\n",
      "                 adversarial Loss : (0.2014 * 0.0500 = 0.0101)\n",
      "                 flownet     Loss : (1.8157 * 2.0000 = 3.6313)\n",
      "                 PSNR  Error      :  16.439236\n",
      "160\n",
      "DiscriminatorModel: Step 160 | Global Loss: 0.453546, lr = 0.000020\n",
      "GeneratorModel : Step 160, lr = 0.000200\n",
      "                 Global      Loss :  3.4193745\n",
      "                 intensity   Loss : (0.1477 * 1.0000 = 0.1477)\n",
      "                 gradient    Loss : (0.1848 * 1.0000 = 0.1848)\n",
      "                 adversarial Loss : (0.1172 * 0.0500 = 0.0059)\n",
      "                 flownet     Loss : (1.5405 * 2.0000 = 3.0810)\n",
      "                 PSNR  Error      :  14.402498\n",
      "170\n",
      "DiscriminatorModel: Step 170 | Global Loss: 0.371937, lr = 0.000020\n",
      "GeneratorModel : Step 170, lr = 0.000200\n",
      "                 Global      Loss :  1.8839625\n",
      "                 intensity   Loss : (0.0722 * 1.0000 = 0.0722)\n",
      "                 gradient    Loss : (0.1285 * 1.0000 = 0.1285)\n",
      "                 adversarial Loss : (0.3086 * 0.0500 = 0.0154)\n",
      "                 flownet     Loss : (0.8339 * 2.0000 = 1.6678)\n",
      "                 PSNR  Error      :  17.492216\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m-------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-6c3ebe810820>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    131\u001b[0m \u001b[0;31m#             print('Training generator...')\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    132\u001b[0m             _, _g_lr, _step, _lp_loss, _gdl_loss, _adv_loss, _flow_loss, _g_loss, _train_psnr, _summaries = sess.run(\n\u001b[0;32m--> 133\u001b[0;31m                 [g_train_op, g_lrate, g_step, lp_loss, gdl_loss, adv_loss, flow_loss, g_loss, train_psnr_error, summary_op])\n\u001b[0m\u001b[1;32m    134\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    135\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0m_step\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;36m10\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Envs/tf/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, fetches, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m    885\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    886\u001b[0m       result = self._run(None, fetches, feed_dict, options_ptr,\n\u001b[0;32m--> 887\u001b[0;31m                          run_metadata_ptr)\n\u001b[0m\u001b[1;32m    888\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mrun_metadata\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    889\u001b[0m         \u001b[0mproto_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf_session\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTF_GetBuffer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrun_metadata_ptr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Envs/tf/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_run\u001b[0;34m(self, handle, fetches, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m   1108\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mfinal_fetches\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mfinal_targets\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mhandle\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mfeed_dict_tensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1109\u001b[0m       results = self._do_run(handle, final_targets, final_fetches,\n\u001b[0;32m-> 1110\u001b[0;31m                              feed_dict_tensor, options, run_metadata)\n\u001b[0m\u001b[1;32m   1111\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1112\u001b[0m       \u001b[0mresults\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Envs/tf/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_do_run\u001b[0;34m(self, handle, target_list, fetch_list, feed_dict, options, run_metadata)\u001b[0m\n\u001b[1;32m   1284\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mhandle\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1285\u001b[0m       return self._do_call(_run_fn, feeds, fetches, targets, options,\n\u001b[0;32m-> 1286\u001b[0;31m                            run_metadata)\n\u001b[0m\u001b[1;32m   1287\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1288\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_do_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_prun_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeeds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetches\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Envs/tf/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_do_call\u001b[0;34m(self, fn, *args)\u001b[0m\n\u001b[1;32m   1290\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_do_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1291\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1292\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1293\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0merrors\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mOpError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1294\u001b[0m       \u001b[0mmessage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompat\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_text\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmessage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Envs/tf/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_run_fn\u001b[0;34m(feed_dict, fetch_list, target_list, options, run_metadata)\u001b[0m\n\u001b[1;32m   1275\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_extend_graph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1276\u001b[0m       return self._call_tf_sessionrun(\n\u001b[0;32m-> 1277\u001b[0;31m           options, feed_dict, fetch_list, target_list, run_metadata)\n\u001b[0m\u001b[1;32m   1278\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1279\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_prun_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeed_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetch_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Envs/tf/lib/python3.6/site-packages/tensorflow/python/client/session.py\u001b[0m in \u001b[0;36m_call_tf_sessionrun\u001b[0;34m(self, options, feed_dict, fetch_list, target_list, run_metadata)\u001b[0m\n\u001b[1;32m   1365\u001b[0m     return tf_session.TF_SessionRun_wrapper(\n\u001b[1;32m   1366\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_session\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeed_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetch_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_list\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1367\u001b[0;31m         run_metadata)\n\u001b[0m\u001b[1;32m   1368\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1369\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_call_tf_sessionprun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeed_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfetch_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# define training generator function\n",
    "with tf.variable_scope('generator', reuse=None):\n",
    "    print('training = {}'.format(tf.get_variable_scope().name))\n",
    "    train_outputs = generator(train_inputs, layers=4, output_channel=3)\n",
    "    train_psnr_error = psnr_error(gen_frames=train_outputs, gt_frames=train_gt)\n",
    "\n",
    "# define testing generator function\n",
    "with tf.variable_scope('generator', reuse=True):\n",
    "    print('testing = {}'.format(tf.get_variable_scope().name))\n",
    "    test_outputs = generator(test_inputs, layers=4, output_channel=3)\n",
    "    test_psnr_error = psnr_error(gen_frames=test_outputs, gt_frames=test_gt)\n",
    "\n",
    "\n",
    "# define intensity loss\n",
    "if lam_lp != 0:\n",
    "    lp_loss = intensity_loss(gen_frames=train_outputs, gt_frames=train_gt, l_num=l_num)\n",
    "else:\n",
    "    lp_loss = tf.constant(0.0, dtype=tf.float32)\n",
    "\n",
    "\n",
    "# define gdl loss\n",
    "if lam_gdl != 0:\n",
    "    gdl_loss = gradient_loss(gen_frames=train_outputs, gt_frames=train_gt, alpha=alpha_num)\n",
    "else:\n",
    "    gdl_loss = tf.constant(0.0, dtype=tf.float32)\n",
    "\n",
    "\n",
    "# define flow loss\n",
    "if lam_flow != 0:\n",
    "    train_gt_flow = flownet(input_a=train_inputs[..., -3:], input_b=train_gt,\n",
    "                            height=flow_height, width=flow_width, reuse=None)\n",
    "    train_pred_flow = flownet(input_a=train_inputs[..., -3:], input_b=train_outputs,\n",
    "                              height=flow_height, width=flow_width, reuse=True)\n",
    "    flow_loss = tf.reduce_mean(tf.abs(train_gt_flow - train_pred_flow))\n",
    "else:\n",
    "    flow_loss = tf.constant(0.0, dtype=tf.float32)\n",
    "\n",
    "\n",
    "# define adversarial loss\n",
    "if adversarial:\n",
    "    with tf.variable_scope('discriminator', reuse=None):\n",
    "        real_logits, real_outputs = discriminator(inputs=train_gt)\n",
    "    with tf.variable_scope('discriminator', reuse=True):\n",
    "        fake_logits, fake_outputs = discriminator(inputs=train_outputs)\n",
    "\n",
    "    print('real_outputs = {}'.format(real_outputs))\n",
    "    print('fake_outputs = {}'.format(fake_outputs))\n",
    "\n",
    "    adv_loss = tf.reduce_mean(tf.square(fake_outputs - 1) / 2)\n",
    "    dis_loss = tf.reduce_mean(tf.square(real_outputs - 1) / 2) + tf.reduce_mean(tf.square(fake_outputs) / 2)\n",
    "else:\n",
    "    adv_loss = tf.constant(0.0, dtype=tf.float32)\n",
    "    dis_loss = tf.constant(0.0, dtype=tf.float32)\n",
    "\n",
    "\n",
    "with tf.name_scope('training'):\n",
    "    g_loss = tf.add_n([lp_loss * lam_lp, gdl_loss * lam_gdl, adv_loss * lam_adv, flow_loss * lam_flow], name='g_loss')\n",
    "\n",
    "    g_step = tf.Variable(0, dtype=tf.int32, trainable=False, name='g_step')\n",
    "    g_lrate = tf.train.piecewise_constant(g_step, boundaries=const.LRATE_G_BOUNDARIES, values=const.LRATE_G)\n",
    "    g_optimizer = tf.train.AdamOptimizer(learning_rate=g_lrate, name='g_optimizer')\n",
    "    g_vars = tf.get_collection(key=tf.GraphKeys.TRAINABLE_VARIABLES, scope='generator')\n",
    "\n",
    "    g_train_op = g_optimizer.minimize(g_loss, global_step=g_step, var_list=g_vars, name='g_train_op')\n",
    "\n",
    "    if adversarial:\n",
    "        # training discriminator\n",
    "        d_step = tf.Variable(0, dtype=tf.int32, trainable=False, name='d_step')\n",
    "        d_lrate = tf.train.piecewise_constant(d_step, boundaries=const.LRATE_D_BOUNDARIES, values=const.LRATE_D)\n",
    "        d_optimizer = tf.train.AdamOptimizer(learning_rate=d_lrate, name='g_optimizer')\n",
    "        d_vars = tf.get_collection(key=tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator')\n",
    "\n",
    "        d_train_op = d_optimizer.minimize(dis_loss, global_step=d_step, var_list=d_vars, name='d_optimizer')\n",
    "    else:\n",
    "        d_step = None\n",
    "        d_lrate = None\n",
    "        d_train_op = None\n",
    "\n",
    "# add all to summaries, for tensorboard\n",
    "tf.summary.scalar(tensor=train_psnr_error, name='train_psnr_error')\n",
    "tf.summary.scalar(tensor=test_psnr_error, name='test_psnr_error')\n",
    "tf.summary.scalar(tensor=g_loss, name='g_loss')\n",
    "tf.summary.scalar(tensor=adv_loss, name='adv_loss')\n",
    "tf.summary.scalar(tensor=dis_loss, name='dis_loss')\n",
    "tf.summary.image(tensor=train_outputs, name='train_outputs')\n",
    "tf.summary.image(tensor=train_gt, name='train_gt')\n",
    "tf.summary.image(tensor=test_outputs, name='test_outputs')\n",
    "tf.summary.image(tensor=test_gt, name='test_gt')\n",
    "summary_op = tf.summary.merge_all()\n",
    "\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = True\n",
    "\n",
    "with tf.Session(config=config) as sess:\n",
    "    # summaries\n",
    "    summary_writer = tf.summary.FileWriter(summary_dir, graph=sess.graph)\n",
    "\n",
    "    # initialize weights\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    print('Init successfully!')\n",
    "\n",
    "    if lam_flow != 0:\n",
    "        # initialize flownet\n",
    "        initialize_flownet(sess, const.FLOWNET_CHECKPOINT)\n",
    "\n",
    "    # tf saver\n",
    "    saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=None)\n",
    "    restore_var = [v for v in tf.global_variables()]\n",
    "    loader = tf.train.Saver(var_list=restore_var)\n",
    "    if os.path.isdir(snapshot_dir):\n",
    "        ckpt = tf.train.get_checkpoint_state(snapshot_dir)\n",
    "        if ckpt and ckpt.model_checkpoint_path:\n",
    "            load(loader, sess, ckpt.model_checkpoint_path)\n",
    "        else:\n",
    "            print('No checkpoint file found.')\n",
    "    else:\n",
    "        load(loader, sess, snapshot_dir)\n",
    "\n",
    "    _step, _loss, _summaries = 0, None, None\n",
    "    while _step < iterations:\n",
    "        try:\n",
    "            \n",
    "            if adversarial:\n",
    "#                 print('Training discriminator...')\n",
    "                _, _d_lr, _d_step, _dis_loss = sess.run([d_train_op, d_lrate, d_step, dis_loss])\n",
    "            else:\n",
    "                _d_step = 0\n",
    "                _d_lr = 0\n",
    "                _dis_loss = 0\n",
    "\n",
    "#             print('Training generator...')\n",
    "            _, _g_lr, _step, _lp_loss, _gdl_loss, _adv_loss, _flow_loss, _g_loss, _train_psnr, _summaries = sess.run(\n",
    "                [g_train_op, g_lrate, g_step, lp_loss, gdl_loss, adv_loss, flow_loss, g_loss, train_psnr_error, summary_op])\n",
    "\n",
    "            if _step % 10 == 0:\n",
    "                print(_step)\n",
    "                print('DiscriminatorModel: Step {} | Global Loss: {:.6f}, lr = {:.6f}'.format(_d_step, _dis_loss, _d_lr))\n",
    "                print('GeneratorModel : Step {}, lr = {:.6f}'.format(_step, _g_lr))\n",
    "                print('                 Global      Loss : ', _g_loss)\n",
    "                print('                 intensity   Loss : ({:.4f} * {:.4f} = {:.4f})'.format(_lp_loss, lam_lp, _lp_loss * lam_lp))\n",
    "                print('                 gradient    Loss : ({:.4f} * {:.4f} = {:.4f})'.format( _gdl_loss, lam_gdl, _gdl_loss * lam_gdl))\n",
    "                print('                 adversarial Loss : ({:.4f} * {:.4f} = {:.4f})'.format(_adv_loss, lam_adv, _adv_loss * lam_adv))\n",
    "                print('                 flownet     Loss : ({:.4f} * {:.4f} = {:.4f})'.format(_flow_loss, lam_flow, _flow_loss * lam_flow))\n",
    "                print('                 PSNR  Error      : ', _train_psnr)\n",
    "            if _step % 100 == 0:\n",
    "                summary_writer.add_summary(_summaries, global_step=_step)\n",
    "                print('Save summaries...')\n",
    "\n",
    "            if _step % 1000 == 0:\n",
    "                save(saver, sess, snapshot_dir, _step)\n",
    "\n",
    "        except tf.errors.OutOfRangeError:\n",
    "            print('Finish successfully!')\n",
    "            save(saver, sess, snapshot_dir, _step)\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adversarial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
